<script setup lang="ts">
import { onMounted, ref } from "vue";
import {
  ElButton,
  ElTree,
  ElIcon,
  ElTreeSelect,
  ElText,
  ElProgress,
} from "element-plus";
import { List } from "@element-plus/icons-vue";
import { getDir } from "@/renderer/api/messageAPI";
import { dir2tree } from "@/renderer/utils/utils";
import { getModelsDir } from "@/renderer/api/machineLearningAPI";
import OCRData from "./OCRData";
import OCRModel, { createOCRModel } from "./OCRModel";
import { NUM_TRAIN_ELEMENTS } from "../datasets/OCRDatasetConstants";

const modelsTree = ref([]);
const datasetsTree = ref([]);
const modelPath = ref();
const datasetPath = ref();
const data = ref();
let model;
const trainProgress = ref(0);

const makeModelsTree = async () => {
  modelsTree.value = [];

  getModelsDir().then((res) => {
    modelsTree.value = dir2tree(res);
    modelsTree.value.forEach((a) =>
      a.children.forEach((b) => (b.disabled = true))
    );
  });
};

const makeDatasetsTree = () => {
  getDir("datasets").then((res) => {
    datasetsTree.value = dir2tree(res);
    datasetsTree.value.forEach((a) =>
      a.children.forEach((b) => (b.disabled = true))
    );
  });
};

const makeOCRModel = async () => {
  await createOCRModel();
  makeModelsTree();
};
// 训练进度条
let betchValue = 0;
const trainProgressFn = (betch, logs) => {
  ++betchValue;

  const percentage = Math.floor(
    ((betchValue * model.BATCH_SIZE) / (NUM_TRAIN_ELEMENTS * model.EPOCHS)) * 100
  );

  trainProgress.value = Math.min(percentage, 100);
};
// 预处理
const preprocessing = async () => {
  if (modelPath.value?.includes("ocr")) {
    model = new OCRModel();

    await model.load();

    model.onBatchEnd = trainProgressFn;
  }

  if (datasetPath.value?.includes("ocr")) {
    data.value = new OCRData();

    await data.value.load();
  }
};

const train = async () => {
  betchValue = 0;

  await model.train(data.value, true);
};

const evaluate = async () => {
  await model.evaluate(data.value);
};

onMounted(() => {
  makeModelsTree();
  makeDatasetsTree();
});
</script>

<template>
  <div class="container">
    <ElButton plain round type="primary" @click="makeOCRModel">
      生成OCR模型
    </ElButton>
    <ElTree class="tree" :data="modelsTree">
      <template #default="{ node }">
        <ElIcon v-if="!node.isLeaf" size="15"><List /></ElIcon>
        <span>{{ node.label }}</span>
      </template>
    </ElTree>
    <ElTreeSelect
      class="tree tree-select"
      v-model="modelPath"
      :data="modelsTree"
      check-strictly
      placeholder="选择模型"
    ></ElTreeSelect>
    <ElTreeSelect
      class="tree tree-select"
      v-model="datasetPath"
      :data="datasetsTree"
      check-strictly
      placeholder="选择数据集"
    ></ElTreeSelect>
    <ElButton plain round type="primary" @click="preprocessing">
      预处理
    </ElButton>
    <ElButton plain round type="primary" @click="train"> 训练 </ElButton>
    <ElProgress
      style="width: 350px"
      :percentage="trainProgress"
      :stroke-width="20"
      :text-inside="true"
      striped
      :striped-flow="trainProgress < 100"
    />
    <ElButton plain round type="primary" @click="evaluate"> 评估 </ElButton>
    <ElText type="warning">样本数: {{ data?.numDatasetElements }}</ElText>
    <ElText type="warning">训练数: {{ data?.numTrainElements }}</ElText>
    <ElText type="warning">验证数: {{ data?.numTestElements }}</ElText>
  </div>
</template>

<style scoped>
.container {
  display: flex;
  flex-direction: column;
}

.container > * {
  margin: 5px;
  align-self: flex-start;
}

.tree {
  max-width: 500px;
  min-width: 300px;
  border-radius: 5px;
  padding: 5px;
}

.tree-select {
  max-width: 200px;
  min-width: 100px;
}

.el-icon {
  margin: 0 5px;
}
</style>
